# -*- coding: utf-8 -*-
"""
删除不要的xml,遮挡图片上的忽略块，空标记样本处理

"""

import os
import sys
import argparse
import cv2
import xml.etree.ElementTree as ET


def file_count(directory):
    return sum([len(files) for root, dirs, files
                in os.walk(directory)])


def del_xml_node(xml_file):
    loop = True
    while loop:
        tree = ET.parse(xml_file)
        root = tree.getroot()
        loop = False
        for obj in root.iter('object'):
            loop = True
            root.remove(obj)
            tree.write(xml_file)


def process(voc_path):
    delete_filelist = open('delete_file.txt', 'w')
    ignore_filelist = open('ignore_file.txt', 'w')
    empty_filelist = open('empty_file.txt', 'w')
    
    xml_path = os.path.join(voc_path, 'Annotations')
    img_path = os.path.join(voc_path, 'JPEGImages')

    count = 0
    # xml_num = file_count(xml_path)
    for root, dirs, files in os.walk(xml_path):
        for file in files:
            delete = False
            ignore = False
            empty = False
            xml_file = os.path.join(xml_path, file)
            tree = ET.parse(xml_file)
            xml_root = tree.getroot()
            ignore_area = []
            for obj in xml_root.iter('object'):
                if obj.find('name').text != 'vehicle':
                    print(obj.find('name').text)
                if obj.find('name').text == 'delete':
                    delete = True
                    delete_filelist.write(file+'\n')
                elif obj.find('name').text == 'ignore':
                    ignore = True
                    xml_box = obj.find('bndbox')
                    xmin = int(xml_box.find('xmin').text)
                    ymin = int(xml_box.find('ymin').text)
                    xmax = int(xml_box.find('xmax').text)
                    ymax = int(xml_box.find('ymax').text)
                    ignore_area.append([xmin, ymin, xmax, ymax])
                    ignore_filelist.write(file+'\n')
                elif obj.find('name').text == 'empty':
                    empty = True
                    empty_filelist.write(file+'\n')
            if delete:
                os.remove(xml_file)
            elif empty:
                del_xml_node(xml_file)
            elif ignore:
                img_path = os.path.join(img_path, file.split('.')[0]+'.jpg')
                img = cv2.imread(img_path)
                for area in ignore_area:
                    img[area[1]:area[3], area[0]:area[2]] = 0
                cv2.imwrite(img_path, img)
            count += 1
            
    delete_filelist.close()
    ignore_filelist.close()
    empty_filelist.close()


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
    process(ROOT_DIR)
